import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from PIL import Image
import torch.nn.functional as F
from model import MyResNet


# ---------------------------------------------------#
#   使用自己训练好的模型预测需要修改model_path参数
# ---------------------------------------------------#
class QZ(object):
    _defaults = {

        "model_path": 'model_data/resnet50-19c8e357.pth',
        # -----------------------------------------------------#
        #   输入图片的大小。
        # -----------------------------------------------------#
        "input_shape": (448, 448, 3),
        # -------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        # -------------------------------#
        "cuda": False
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    # ---------------------------------------------------#
    #   初始化Siamese
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        self.generate()

    # ---------------------------------------------------#
    #   载入模型
    # ---------------------------------------------------#
    def generate(self):
        # ---------------------------#
        #   载入模型与权值
        # ---------------------------#
        # print('Loading weights into state dict...')
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = MyResNet()
        model.load_state_dict(torch.load(self.model_path, map_location=device), False)
        self.net = model.eval()
        # print('{} model loaded.'.format(self.model_path))

        if self.cuda:
            self.net = torch.nn.DataParallel(self.net)#多卡
            cudnn.benchmark = True
            self.net = self.net.cuda()

    def letterbox_image(self, image, size):
        image = image.convert("RGB")
        iw, ih = image.size
        w, h = size
        scale = min(w / iw, h / ih)
        nw = int(iw * scale)
        nh = int(ih * scale)

        image = image.resize((nw, nh), Image.BICUBIC)
        new_image = Image.new('RGB', size, (128, 128, 128))
        new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
        if self.input_shape[-1] == 1:
            new_image = new_image.convert("L")
        return new_image

    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image_1):
        # ---------------------------------------------------#
        #   对输入图像进行不失真的resize
        # ---------------------------------------------------#
        image_1 = self.letterbox_image(image_1, [self.input_shape[1], self.input_shape[0]])

        # ---------------------------------------------------#
        #   对输入图像进行归一化
        # ---------------------------------------------------#
        photo_1 = np.asarray(image_1).astype(np.float64) / 255

        if self.input_shape[-1] == 1:
            photo_1 = np.expand_dims(photo_1, -1)

        with torch.no_grad():
            # ---------------------------------------------------#
            #   添加上batch维度，才可以放入网络中预测
            # ---------------------------------------------------#
            photo_1 = torch.from_numpy(np.expand_dims(np.transpose(photo_1, (2, 0, 1)), 0)).type(torch.FloatTensor)

            photo_1 = F.interpolate(photo_1, size=(224, 224), mode='nearest')

            if self.cuda:
                photo_1 = photo_1.cuda()

            # ---------------------------------------------------#
            #   获得预测结果，output输出为概率
            # ---------------------------------------------------#

            output = self.net(photo_1)[0]
            output = torch.nn.Sigmoid()(output)
            _, preds = torch.max(output, 1)
            # output = preds.view(-1)

        return preds[0].item()

def api(path):
    model = QZ()
    image = Image.open(path)
    probability = model.detect_image(image)
    if probability == 1:
        print("这是重载")
    else:
        print("这是轻载")

if __name__ == "__main__":
    path = "../ship/val/Z/192.168.101.244_01_20211204155028607.jpg"
    api(path)